#include <mcp/server/fastmcp.h>
#include <iostream>
#include <string>
#include <vector>
#include <cmath>
#include <memory>

using namespace mcp;
using namespace mcp::server;
using namespace mcp::types;

// Basic arithmetic operations
Result<json::Value> add_tool(const json::Value& args, Context<>& context) {
    auto a = args["a"].as_double();
    auto b = args["b"].as_double();
    double result = a + b;

    context.info(std::format("Addition: {} + {} = {}", a, b, result));
    context.report_progress(1.0, 1.0, "Addition complete");

    return json::Value::object({
        {"operation", "add"},
        {"operand1", a},
        {"operand2", b},
        {"result", result}
    });
}

Result<json::Value> subtract_tool(const json::Value& args, Context<>& context) {
    auto a = args["a"].as_double();
    auto b = args["b"].as_double();
    double result = a - b;

    context.info(std::format("Subtraction: {} - {} = {}", a, b, result));
    context.report_progress(1.0, 1.0, "Subtraction complete");

    return json::Value::object({
        {"operation", "subtract"},
        {"operand1", a},
        {"operand2", b},
        {"result", result}
    });
}

Result<json::Value> multiply_tool(const json::Value& args, Context<>& context) {
    auto a = args["a"].as_double();
    auto b = args["b"].as_double();
    double result = a * b;

    context.info(std::format("Multiplication: {} × {} = {}", a, b, result));
    context.report_progress(1.0, 1.0, "Multiplication complete");

    return json::Value::object({
        {"operation", "multiply"},
        {"operand1", a},
        {"operand2", b},
        {"result", result}
    });
}

Result<json::Value> divide_tool(const json::Value& args, Context<>& context) {
    auto a = args["a"].as_double();
    auto b = args["b"].as_double();

    if (b == 0.0) {
        return Error::invalid_params("Division by zero is not allowed");
    }

    double result = a / b;

    context.info(std::format("Division: {} ÷ {} = {}", a, b, result));
    context.report_progress(1.0, 1.0, "Division complete");

    return json::Value::object({
        {"operation", "divide"},
        {"operand1", a},
        {"operand2", b},
        {"result", result}
    });
}

// Advanced math operations
Result<json::Value> power_tool(const json::Value& args, Context<>& context) {
    auto base = args["base"].as_double();
    auto exponent = args["exponent"].as_double();
    double result = std::pow(base, exponent);

    context.info(std::format("Power: {}^{} = {}", base, exponent, result));
    context.report_progress(1.0, 1.0, "Power calculation complete");

    return json::Value::object({
        {"operation", "power"},
        {"base", base},
        {"exponent", exponent},
        {"result", result}
    });
}

Result<json::Value> sqrt_tool(const json::Value& args, Context<>& context) {
    auto value = args["value"].as_double();

    if (value < 0.0) {
        return Error::invalid_params("Cannot calculate square root of negative number");
    }

    double result = std::sqrt(value);

    context.info(std::format("Square root: √{} = {}", value, result));
    context.report_progress(1.0, 1.0, "Square root complete");

    return json::Value::object({
        {"operation", "sqrt"},
        {"value", value},
        {"result", result}
    });
}

// Statistics tool
Result<json::Value> statistics_tool(const json::Value& args, Context<>& context) {
    auto numbers = args["numbers"].as_array();
    if (numbers.empty()) {
        return Error::invalid_params("Numbers array cannot be empty");
    }

    double sum = 0.0;
    double min = numbers[0].as_double();
    double max = numbers[0].as_double();

    for (const auto& num : numbers) {
        double value = num.as_double();
        sum += value;
        min = std::min(min, value);
        max = std::max(max, value);
    }

    double mean = sum / numbers.size();
    double variance = 0.0;

    for (const auto& num : numbers) {
        double diff = num.as_double() - mean;
        variance += diff * diff;
    }
    variance /= numbers.size();

    double std_dev = std::sqrt(variance);

    context.info(std::format("Statistics calculated for {} numbers", numbers.size()));
    context.report_progress(1.0, 1.0, "Statistics complete");

    return json::Value::object({
        {"operation", "statistics"},
        {"count", static_cast<int64_t>(numbers.size())},
        {"sum", sum},
        {"mean", mean},
        {"min", min},
        {"max", max},
        {"variance", variance},
        {"standard_deviation", std_dev}
    });
}

// Math constants resource
Result<ResourceContent> constants_resource(const std::string& uri, Context<>& context) {
    context.info("Math constants resource requested");

    std::string content = R"(
Mathematical Constants
====================

π (pi) = 3.14159265358979323846
e (Euler's number) = 2.71828182845904523536
φ (golden ratio) = 1.61803398874989484820
√2 = 1.41421356237309504880
γ (Euler-Mascheroni) = 0.57721566490153286061

Common Formulas:
- Circle area = π × r²
- Circle circumference = 2 × π × r
- Sphere volume = (4/3) × π × r³
- Pythagorean theorem: a² + b² = c²
)";

    return ResourceContent{
        .contents = std::vector<std::shared_ptr<ContentBlock>>{
            std::make_shared<TextContent>(content)
        },
        .uri = uri
    };
}

// Math problem prompt
Result<std::vector<PromptMessage>> math_problem_prompt(const json::Value& args, Context<>& context) {
    std::string difficulty = "medium";
    std::string topic = "arithmetic";

    if (args.contains("difficulty")) {
        difficulty = args["difficulty"].as_string();
    }
    if (args.contains("topic")) {
        topic = args["topic"].as_string();
    }

    std::string problem;
    if (topic == "arithmetic") {
        if (difficulty == "easy") {
            problem = "What is 15 + 27?";
        } else if (difficulty == "medium") {
            problem = "A rectangle has a length of 12 cm and width of 8 cm. What is its area?";
        } else {
            problem = "Solve for x: 3x + 7 = 22";
        }
    } else if (topic == "geometry") {
        if (difficulty == "easy") {
            problem = "What is the area of a circle with radius 5 cm?";
        } else if (difficulty == "medium") {
            problem = "Find the hypotenuse of a right triangle with sides 6 cm and 8 cm.";
        } else {
            problem = "Calculate the volume of a sphere with radius 4 cm.";
        }
    } else {
        problem = "Create an interesting math problem related to " + topic + " at " + difficulty + " difficulty.";
    }

    context.info(std::format("Generated {} {} math problem", difficulty, topic));

    return std::vector<PromptMessage>{
        PromptMessage{
            .role = PromptRole::System,
            .content = std::make_shared<TextContent>(
                "You are a math tutor. Solve the problem step by step and explain your reasoning clearly."
            )
        },
        PromptMessage{
            .role = PromptRole::User,
            .content = std::make_shared<TextContent>(problem)
        }
    };
}

int main() {
    try {
        std::cout << "MCP C++ Math Server Example\n";
        std::cout << "===========================\n\n";

        // Create server configuration
        ServerConfig config{
            .name = "MCP Math Server",
            .version = "1.0.0",
            .description = "Mathematical operations server with various tools",
            .debug = true,
            .host = "localhost",
            .port = 8082
        };

        // Create server
        auto server = std::make_shared<FastMCPServer>(config);

        // Register basic arithmetic tools
        server->tool("add", add_tool, ToolOptions{
            .description = "Add two numbers",
            .input_schema = json::object_schema()
                .property("a", json::number_schema().description("First number"))
                .property("b", json::number_schema().description("Second number"))
                .required({"a", "b"})
        });

        server->tool("subtract", subtract_tool, ToolOptions{
            .description = "Subtract second number from first",
            .input_schema = json::object_schema()
                .property("a", json::number_schema().description("First number"))
                .property("b", json::number_schema().description("Second number"))
                .required({"a", "b"})
        });

        server->tool("multiply", multiply_tool, ToolOptions{
            .description = "Multiply two numbers",
            .input_schema = json::object_schema()
                .property("a", json::number_schema().description("First number"))
                .property("b", json::number_schema().description("Second number"))
                .required({"a", "b"})
        });

        server->tool("divide", divide_tool, ToolOptions{
            .description = "Divide first number by second",
            .input_schema = json::object_schema()
                .property("a", json::number_schema().description("First number"))
                .property("b", json::number_schema().description("Second number"))
                .required({"a", "b"})
        });

        // Register advanced math tools
        server->tool("power", power_tool, ToolOptions{
            .description = "Calculate base raised to exponent",
            .input_schema = json::object_schema()
                .property("base", json::number_schema().description("Base number"))
                .property("exponent", json::number_schema().description("Exponent"))
                .required({"base", "exponent"})
        });

        server->tool("sqrt", sqrt_tool, ToolOptions{
            .description = "Calculate square root",
            .input_schema = json::object_schema()
                .property("value", json::number_schema().description("Number to calculate square root of"))
                .required({"value"})
        });

        server->tool("statistics", statistics_tool, ToolOptions{
            .description = "Calculate statistics for an array of numbers",
            .input_schema = json::object_schema()
                .property("numbers", json::array_schema()
                    .items(json::number_schema())
                    .description("Array of numbers"))
                .required({"numbers"})
        });

        // Register resource
        server->resource("math://constants", constants_resource, ResourceOptions{
            .name = "Math Constants",
            .description = "Common mathematical constants and formulas",
            .mime_type = "text/plain"
        });

        // Register prompt
        server->prompt("math_problem", math_problem_prompt, PromptOptions{
            .description = "Generate a math problem"
        });

        // Display server info
        std::cout << "Server Info:\n";
        std::cout << "  Name: " << config.name << "\n";
        std::cout << "  Version: " << config.version << "\n";
        std::cout << "  Port: " << config.port << "\n\n";

        // Display registered tools
        auto tools = server->list_tools();
        std::cout << "Available Tools (" << tools.size() << "):\n";
        for (const auto& tool : tools) {
            std::cout << "  - " << tool.name << ": " << tool.description << "\n";
        }

        std::cout << "\nAvailable Resources (" << server->list_resources().size() << "):\n";
        for (const auto& resource : server->list_resources()) {
            std::cout << "  - " << resource.uri << ": " << resource.name << "\n";
        }

        std::cout << "\nAvailable Prompts (" << server->list_prompts().size() << "):\n";
        for (const auto& prompt : server->list_prompts()) {
            std::cout << "  - " << prompt.name << ": "
                     << (prompt.description ? *prompt.description : "No description") << "\n";
        }

        std::cout << "\nStarting server with stdio transport...\n";
        std::cout << "Example usage:\n";
        std::cout << "  {\"method\":\"tools/call\",\"params\":{\"name\":\"add\",\"arguments\":{\"a\":15,\"b\":25}}}\n";
        std::cout << "  {\"method\":\"tools/call\",\"params\":{\"name\":\"statistics\",\"arguments\":{\"numbers\":[1,2,3,4,5]}}}\n";
        std::cout << "  {\"method\":\"resources/read\",\"params\":{\"uri\":\"math://constants\"}}\n";
        std::cout << "\nPress Ctrl+C to stop the server\n";

        // Run server
        server->run<mcp::transport::StdioTransport>();

    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << "\n";
        return 1;
    }

    return 0;
}